package step

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"github.com/wanchain/go-mpc/common"
	"github.com/wanchain/go-mpc/common/hexutil"
	"github.com/wanchain/go-mpc/log"
	"github.com/wanchain/go-mpc/storeman/osmconf"
	schcomm "github.com/wanchain/go-mpc/storeman/schnorrcomm"
	mpcprotocol "github.com/wanchain/go-mpc/storeman/storemanmpc/protocol"
	"math/big"
	"strconv"
)

type MpcRStepUni struct {
	BaseMpcStep
	resultKeys []string
	signNum    int

	RShareErrNum    int
	rpkshareOKIndex []uint16
	rpkshareKOIndex []uint16
	rpkshareNOIndex []uint16
	sigNum          int
}

func CreateMpcRStepUni(totalNum int, sigNum int) *MpcRStepUni {
	resultkeys := []string{mpcprotocol.RPk}
	mpc := &MpcRStepUni{*CreateBaseMpcStep(totalNum, sigNum),
		resultkeys,
		sigNum,
		0,
		make([]uint16, 0),
		make([]uint16, 0),
		make([]uint16, 0),
		sigNum}

	for i := 0; i < sigNum; i++ {
		mpc.messages[i] = createPointGeneratorUni(i)
	}
	return mpc
}

func (addStep *MpcRStepUni) CreateMessage() []mpcprotocol.StepMessage {
	log.SyslogInfo("MpcRStepUni::CreateMessage begin", "ctx", addStep.GetMpcContextId())
	message := make([]mpcprotocol.StepMessage, 1)
	message[0].MsgCode = mpcprotocol.MPCMessage
	message[0].PeerID = nil

	for sigIndex := 0; sigIndex < addStep.signNum; sigIndex++ {

		pointer := addStep.messages[sigIndex].(*mpcPointGeneratorUni)
		seedBytes, err := addStep.schnorrMpcer.MarshPt(pointer.seed)
		if err != nil {
			log.SyslogErr("MpcRStepUni CreateMessage", "MarshPt err", err.Error(), "ctx", addStep.GetMpcContextId())
		}

		// only one point, rpkShare
		message[0].BytesData = append(message[0].BytesData, seedBytes)
	}

	var buf bytes.Buffer
	for _, value := range message[0].BytesData {
		buf.Write(value)
	}
	h := sha256.Sum256(buf.Bytes())

	prv, err := osmconf.GetOsmConf().GetSelfPrvKey()
	if err != nil {
		return nil
	}
	r, s, _ := schcomm.SignInternalData(prv, h[:])
	message[0].Data = make([]big.Int, 2)
	message[0].Data[0] = *r
	message[0].Data[1] = *s

	return message
}

func (addStep *MpcRStepUni) CheckRCommit(nodeIndex uint16, RCommitComputed common.Hash) bool {

	key := mpcprotocol.RCommit + strconv.Itoa(int(nodeIndex))

	RCommitSaved, err := addStep.mpcResult.GetByteValue(key)

	if err != nil {
		return false
	}
	log.SyslogDebug("MpcRStepUni::CheckRCommit begin ",
		"RCommitSaved", common.ToHex(RCommitSaved),
		"RCommitComputed ", RCommitComputed.Hex(), "nodeIndex", nodeIndex, "ctx", addStep.GetMpcContextId())

	return bytes.Equal(RCommitSaved[:], RCommitComputed[:])

}

func (addStep *MpcRStepUni) HandleMessage(msg *mpcprotocol.StepMessage) bool {

	log.SyslogInfo("MpcRStepUni::HandleMessage begin ",
		"peerID", msg.PeerID.SlimString(),
		"gpk ", hex.EncodeToString(msg.BytesData[0]), "ctx", addStep.GetMpcContextId())

	_, grpIdString, _ := addStep.BaseStep.GetGrpId()

	senderPk, _ := osmconf.GetOsmConf().GetPKByNodeId(grpIdString, msg.PeerID)

	err := schcomm.CheckPK(senderPk)
	if err != nil {
		log.SyslogErr("MpcRStepUni", "HandleMessage", err.Error(), "ctx", addStep.GetMpcContextId())
	}

	senderIndex, _ := osmconf.GetOsmConf().GetInxByNodeId(grpIdString, msg.PeerID)
	var buf bytes.Buffer
	for _, value := range msg.BytesData {
		buf.Write(value)
	}
	h := sha256.Sum256(buf.Bytes())
	r := msg.Data[0]
	s := msg.Data[1]

	// add check RCommit begin
	RCommitComputed := sha256.Sum256(buf.Bytes())
	if !addStep.CheckRCommit(senderIndex, RCommitComputed) {
		log.SyslogErr("MpcRStepUni::HandleMessage", "MpcRStepUni.CheckRCommit failed", msg.PeerID.SlimString(), "ctx", addStep.GetMpcContextId())
		return false
	}
	// add check RCommit end
	bVerifySig := schcomm.VerifyInternalData(senderPk, h[:], &r, &s)
	if bVerifySig {

		for sigIndex := 0; sigIndex < addStep.signNum; sigIndex++ {
			pointer := addStep.messages[sigIndex].(*mpcPointGeneratorUni)
			_, exist := pointer.message[*msg.PeerID]
			if exist {
				log.SyslogErr("MpcRStepUni::HandleMessage", "MpcPointStep.HandleMessage, get msg from seed fail. peer", msg.PeerID.SlimString(), "ctx", addStep.GetMpcContextId())
				return false
			}

			pointPk, err := addStep.schnorrMpcer.UnMarshPt(msg.BytesData[sigIndex])
			if err != nil {
				log.SyslogErr("MpcRStepUni::HandleMessage", "UnMarshPt error ", err.Error(), "ctx", addStep.GetMpcContextId())
				return false
			}

			pointer.message[*msg.PeerID] = pointPk
			// save rpkshare for check data of s
			key := mpcprotocol.RPkShare + strconv.Itoa(int(senderIndex)) + strconv.Itoa(sigIndex)
			addStep.mpcResult.SetByteValue(key, msg.BytesData[sigIndex])
			log.SyslogInfo("MpcRStepUni::HandleMessage save rpkshare", "key", key, "rpkshare", hexutil.Encode(msg.BytesData[0]), "ctx", addStep.GetMpcContextId())
		}

		addStep.rpkshareOKIndex = append(addStep.rpkshareOKIndex, uint16(senderIndex))
	} else {
		log.SyslogErr("MpcPointStep::HandleMessage", " check sig fail", "ctx", addStep.GetMpcContextId())
		addStep.rpkshareKOIndex = append(addStep.rpkshareKOIndex, uint16(senderIndex))
	}

	return true
}

func (addStep *MpcRStepUni) FinishStep(result mpcprotocol.MpcResultInterface, mpc mpcprotocol.StoremanManager) error {
	err := addStep.BaseMpcStep.FinishStep()
	if err != nil {
		return err
	}

	for sigIndex := 0; sigIndex < addStep.signNum; sigIndex++ {
		pointer := addStep.messages[sigIndex].(*mpcPointGeneratorUni)
		resultBytes, err := addStep.schnorrMpcer.MarshPt(pointer.result)
		if err != nil {
			log.SyslogErr("MpcRStepUni FinishStep", "MarshPt error", err.Error(), "ctx", addStep.GetMpcContextId())
			return err
		}
		log.Info("MpcRStepUni::FinishStep",
			"result key", addStep.resultKeys[0],
			"result value ", hexutil.Encode(resultBytes), "ctx", addStep.GetMpcContextId())

		key := mpcprotocol.RPk + strconv.Itoa(sigIndex)
		err = result.SetByteValue(key, resultBytes)
		if err != nil {
			log.SyslogErr("MpcRStepUni::FinishStep HandleMessage", "MpcPointStep.FinishStep, SetValue fail. err", err.Error(), "ctx", addStep.GetMpcContextId())
			return err
		}
		log.SyslogInfo("MpcRStepUni::FinishStep save RPK", "key", key, "RPK", hexutil.Encode(resultBytes), "ctx", addStep.GetMpcContextId())
	}

	_, grpIdString, _ := addStep.BaseStep.GetGrpId()

	allIndex, _ := osmconf.GetOsmConf().GetGrpElemsInxes(grpIdString)
	tempIndex := osmconf.Difference(*allIndex, addStep.rpkshareOKIndex)
	addStep.rpkshareNOIndex = osmconf.Difference(tempIndex, addStep.rpkshareKOIndex)

	okIndex := make([]big.Int, 0)
	koIndex := make([]big.Int, 0)
	noIndex := make([]big.Int, 0)

	for _, value := range addStep.rpkshareOKIndex {
		okIndex = append(okIndex, *big.NewInt(0).SetInt64(int64(value)))
		//okIndex[i].SetInt64(int64(value))
	}

	for _, value := range addStep.rpkshareKOIndex {
		koIndex = append(koIndex, *big.NewInt(0).SetInt64(int64(value)))
		//koIndex[i].SetInt64(int64(value))
	}

	for _, value := range addStep.rpkshareNOIndex {
		//noIndex[i].SetInt64(int64(value))
		noIndex = append(noIndex, *big.NewInt(0).SetInt64(int64(value)))
	}

	addStep.mpcResult.SetValue(mpcprotocol.ROKIndex, okIndex)
	addStep.mpcResult.SetValue(mpcprotocol.RKOIndex, koIndex)
	addStep.mpcResult.SetValue(mpcprotocol.RNOIndex, noIndex)

	okBig, _ := osmconf.BuildDataByIndexes(&okIndex)
	koBig, _ := osmconf.BuildDataByIndexes(&koIndex)
	noBig, _ := osmconf.BuildDataByIndexes(&noIndex)
	log.SyslogInfo("MpcRStepUni::FinishStep",
		"rpkshareOKIndex", hex.EncodeToString(okBig.Bytes()),
		"rpkshareKOindex", hex.EncodeToString(koBig.Bytes()),
		"rpkshareNOindex", hex.EncodeToString(noBig.Bytes()),
		"ctx", addStep.GetMpcContextId())

	if err != nil {
		_, retHash := addStep.BaseMpcStep.GetSignedDataHash(result)
		addStep.BaseMpcStep.ShowNotArriveNodes(retHash, mpc.SelfNodeId())
		return err
	}
	return nil
}
